import os
import requests
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
import subprocess

# 下载加密的 m3u8 文件
def download_m3u8(url, output_path):
    response = requests.get(url)
    with open(output_path, 'wb') as f:
        f.write(response.content)

# 解析 m3u8 文件并下载加密的 ts 文件
def parse_m3u8(m3u8_path, ts_output_dir):
    with open(m3u8_path, 'r') as f:
        lines = f.readlines()
    ts_urls = []
    for line in lines:
        if line.startswith('#'):
            continue
        ts_urls.append(line.strip())
    if not os.path.exists(ts_output_dir):
        os.makedirs(ts_output_dir)
    for i, ts_url in enumerate(ts_urls):
        response = requests.get(ts_url)
        with open(os.path.join(ts_output_dir, f'{i}.ts'), 'wb') as f:
            f.write(response.content)

# 解密 ts 文件
def decrypt_ts(ts_dir, key, key_id, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for file in os.listdir(ts_dir):
        if file.endswith('.ts'):
            file_path = os.path.join(ts_dir, file)
            with open(file_path, 'rb') as f:
                encrypted_data = f.read()
            # 提取 IV（通常在 m3u8 中指定，这里假设 IV 是 16 个字节的 0）
            iv = bytes.fromhex('00000000000000000000000000000000')
            cipher = AES.new(key, AES.MODE_CBC, iv)
            decrypted_data = cipher.decrypt(encrypted_data)
            decrypted_data = unpad(decrypted_data, AES.block_size)
            output_path = os.path.join(output_dir, file)
            with open(output_path, 'wb') as f:
                f.write(decrypted_data)

# 使用 FFmpeg 将解密后的 ts 文件转码为 HLS
def transcode_to_hls(decrypted_ts_dir, output_dir):
    # 创建一个包含所有解密后 ts 文件的列表文件
    with open(os.path.join(decrypted_ts_dir, 'file_list.txt'), 'w') as f:
        for file in sorted(os.listdir(decrypted_ts_dir)):
            if file.endswith('.ts'):
                f.write(f"file '{os.path.join(decrypted_ts_dir, file)}'\n")
    # 调用 FFmpeg 进行转码
    cmd = [
        'ffmpeg',
        '-f', 'concat',
        '-i', os.path.join(decrypted_ts_dir, 'file_list.txt'),
        '-c:v', 'libx264',
        '-preset', 'veryfast',
        '-maxrate', '3000k',
        '-bufsize', '6000k',
        '-pix_fmt', 'yuv420p',
        '-g', '50',
        '-c:a', 'aac',
        '-b:a', '160k',
        '-ac', '2',
        '-ar', '44100',
        '-f', 'hls',
        '-hls_time', '10',
        '-hls_list_size', '6',
        '-hls_flags', 'delete_segments+append_list',
        '-hls_segment_filename', os.path.join(output_dir, '%05d.ts'),
        os.path.join(output_dir, 'playlist.m3u8')
    ]
    subprocess.run(cmd)

# 示例用法
if __name__ == "__main__":
    # 输入加密的 m3u8 URL
    encrypted_m3u8_url = 'http://62.72.163.96:5000/mytvsuper/J'
    # ClearKey 的密钥和密钥 ID（需要从 DRM 系统获取）
    key = bytes.fromhex('2d2fd7b1661b1e28de38268872b48480')  # 替换为实际的 16 字节密钥（32 字符的十六进制字符串）
    key_id = '0958b9c657622c465a6205eb2252b8ed'  # 替换为实际的密钥 ID

    # 下载并解析加密的 m3u8 文件
    m3u8_output_path = 'encrypted_stream.m3u8'
    download_m3u8(encrypted_m3u8_url, m3u8_output_path)
    ts_output_dir = 'encrypted_ts'
    parse_m3u8(m3u8_output_path, ts_output_dir)

    # 解密 ts 文件
    decrypted_ts_dir = 'decrypted_ts'
    decrypt_ts(ts_output_dir, key, key_id, decrypted_ts_dir)

    # 转码为 HLS
    hls_output_dir = 'hls_output'
    transcode_to_hls(decrypted_ts_dir, hls_output_dir)